# Bootstrap main model results
load(working.data.loc)
# B = 1000
# B = 3
# B = 500
B = 200

lm.regs.backup = copy(lm.regs) # save backup

# coeftest(lm.onshore.oil, vcov.=NeweyWest)
# Functions to modify lm model fit
sample.lm.coefs = function(lmfit, B) {
  beta = lmfit$coefficients
  V.temp = NeweyWest(lmfit)
  beta.sample = MASS::mvrnorm(B, mu=beta, Sigma=V.temp)
  beta.sample = t(beta.sample)
  return(beta.sample)
}
modify.lmfit = function(lmfit, beta) {
  lmcopy = copy(lmfit)
  lmcopy$coefficients = beta
  return(lmcopy)
}
sample.sys.coefs = function(sysfit, V.temp, B) {
  beta = sysfit$coefficients
  
  beta.sample = MASS::mvrnorm(B, mu=beta, Sigma=V.temp)
  beta.sample = t(beta.sample)
  return(beta.sample)
}
modify.lmfit.for.sys = function(lmfit, beta, regex) {
  # extract relevant beta coefficients
  beta = beta[grep(regex, names(beta))]
  lmcopy = copy(lmfit)
  if (any(gsub(regex,'', names(beta))!=names(lmcopy$coefficients))) message('Coefficients do not align')
  lmcopy$coefficients = beta
  return(lmcopy)
}
# Sample parameters
library(systemfit)
V.temp = NeweyWest(sur.fit)
set.seed(2021) 
beta.sample = sample.sys.coefs(sur.fit, V.temp, B)
apply(beta.sample, 1, mean)/sur.fit$coefficients
hist(cor(t(beta.sample))-cov2cor(V.temp), breaks=1000)
mean(cor(t(beta.sample))-cov2cor(V.temp))

tot.effects.sample = matrix(NA, nrow=6, ncol=B)
tot.effects.sample[1,] = apply(beta.sample, 2, function(x) sum(x[grep('eq1.*wti', rownames(beta.sample))]))
tot.effects.sample[2,] = apply(beta.sample, 2, function(x) sum(x[grep('eq2.*wti', rownames(beta.sample))]))
tot.effects.sample[3,] = apply(beta.sample, 2, function(x) sum(x[grep('eq3.*wti', rownames(beta.sample))]))
tot.effects.sample[4,] = apply(beta.sample, 2, function(x) sum(x[grep('eq1.*hh', rownames(beta.sample))]))
tot.effects.sample[5,] = apply(beta.sample, 2, function(x) sum(x[grep('eq2.*hh', rownames(beta.sample))]))
tot.effects.sample[6,] = apply(beta.sample, 2, function(x) sum(x[grep('eq3.*hh', rownames(beta.sample))]))
par(mfrow=c(2,3))
hist(tot.effects.sample[1,], main='Oil Wells, WTI')
hist(tot.effects.sample[2,], main='Gas Wells, WTI')
hist(tot.effects.sample[3,], main='Offshore Wells, WTI')
hist(tot.effects.sample[4,], main='Oil Wells, HH')
hist(tot.effects.sample[5,], main='Gas Wells, HH')
hist(tot.effects.sample[6,], main='Offshore Wells, HH')
apply(tot.effects.sample, 1, function(x) mean(x<0)) # in terms of own price elasticities, we're ok. can't throw away negatives since that would entail throwing away 8-20% of offshore draws.
cor(t(tot.effects.sample))

# beta.onshore.oil.sample = sample.lm.coefs(lm.onshore.oil, B)
# beta.onshore.gas.sample = sample.lm.coefs(lm.onshore.gas, B)
# beta.offshore.sample = sample.lm.coefs(lm.offshore, B)

# ## Check that sampled distribution looks like mean and variance
# apply(beta.onshore.oil.sample, 1, mean)/lm.onshore.oil$coefficients
# hist(cor(t(beta.onshore.oil.sample))-cov2cor(NeweyWest(lm.onshore.oil)), breaks=1000); axis(1, at=1)
# mean(cor(t(beta.onshore.oil.sample))-cov2cor(NeweyWest(lm.onshore.oil)))
# 
# apply(beta.onshore.gas.sample, 1, mean)/lm.onshore.gas$coefficients
# hist(cor(t(beta.onshore.gas.sample))-cov2cor(NeweyWest(lm.onshore.gas)), breaks=1000); axis(1, at=1)
# mean(cor(t(beta.onshore.gas.sample))-cov2cor(NeweyWest(lm.onshore.gas)))
# cov(t(beta.onshore.gas.sample))[which(cov(t(beta.onshore.gas.sample))/NeweyWest(lm.onshore.gas)>15)]
# 
# apply(beta.offshore.sample, 1, mean)/lm.offshore$coefficients
# hist(cor(t(beta.offshore.sample))-cov2cor(NeweyWest(lm.offshore)), breaks=1000); axis(1, at=1)
# mean(cor(t(beta.offshore.sample))-cov2cor(NeweyWest(lm.offshore)))
# cov(t(beta.offshore.sample))[which(cov(t(beta.offshore.sample))/NeweyWest(lm.offshore)>15)]

bootstrap.results = array(NA, dim=list(2,8,B))
bootstrap.results = array(NA, dim=list(6,8,B))
library(progress)
pb <- progress_bar$new(format = " Running bootstrap sims [:bar] :percent ETA= :eta",
                       total = B, clear = FALSE, width= 60)
st = Sys.time()
for (b in 1:B) {
  ## Modify lm fits
  # lm.onshore.oil.sample = modify.lmfit(lm.onshore.oil, beta.onshore.oil.sample[,b])
  # lm.onshore.gas.sample = modify.lmfit(lm.onshore.gas, beta.onshore.gas.sample[,b])
  # lm.offshore.sample = modify.lmfit(lm.offshore, beta.offshore.sample[,b])
  
  ## Modify lm fits based on joint covariance matrix
  lm.onshore.oil.sample = modify.lmfit.for.sys(lm.onshore.oil, beta.sample[,b], regex='eq1_')
  lm.onshore.gas.sample = modify.lmfit.for.sys(lm.onshore.gas, beta.sample[,b], regex='eq2_')
  lm.offshore.sample = modify.lmfit.for.sys(lm.offshore, beta.sample[,b], regex='eq3_')
  
  # Store
  regs.use.sample = array(list(), dim=dim(lm.regs), dimnames=dimnames(lm.regs))
  regs.use.sample['Oil',,'Onshore'] = list(lm.onshore.oil.sample, lm.onshore.oil.sample)
  regs.use.sample['Gas',,'Onshore'] = list(lm.onshore.gas.sample, lm.onshore.gas.sample)
  regs.use.sample[,,'Offshore'] = list(lm.offshore.sample, lm.offshore.sample, lm.offshore.sample, lm.offshore.sample)
  
  # Overwrite operative lm fits
  regs.use = regs.use.sample
  # Recalibrate
  source(script.dir%&%'calibrate_model.R')
  
  # Simulate baseline, $20/ton adder, and moratorium
  sim.base.bs = simulate.policy()
  sim.18.75.royalty.rate.on.bs = simulate.policy(roy.fed.new.on.sim=0.1875)
  sim.25.royalty.rate.on.bs = simulate.policy(roy.fed.new.on.sim=0.25)
  sim.25.royalty.rate.bs = simulate.policy(roy.fed.new.on.sim=0.25, roy.fed.new.off.sim=0.25)
  sim.carbon.adder.fed.bs = simulate.policy(carb.add.fed.sim=50)
  sim.carbon.adder.fed.20.bs = simulate.policy(carb.add.fed.sim=20)
  sim.ban.lag.0.bs = simulate.policy(fed.leasing.ban.sim = TRUE, ban.lag=0)
  
  # Output key results
  sim.list.bs = list(sim.base=sim.base.bs$sim,
                     sim.18.75.royalty.rate.on=sim.18.75.royalty.rate.on.bs$sim,
                     sim.25.royalty.rate.on=sim.25.royalty.rate.on.bs$sim,
                     sim.25.royalty.rate=sim.25.royalty.rate.bs$sim,
                     sim.carbon.adder.fed.20=sim.carbon.adder.fed.20.bs$sim,
                     sim.carbon.adder.fed=sim.carbon.adder.fed.bs$sim,
                     sim.ban.lag.0=sim.ban.lag.0.bs$sim)
  
  sim.name.list.bs = names(sim.list.bs)
  
  mr.list.bs = list(sim.base=sim.base.bs$mr,
                    sim.18.75.royalty.rate.on=sim.18.75.royalty.rate.on.bs$mr,
                    sim.25.royalty.rate.on=sim.25.royalty.rate.on.bs$mr,
                    sim.25.royalty.rate=sim.25.royalty.rate.bs$mr,
                    sim.carbon.adder.fed.20=sim.carbon.adder.fed.20.bs$mr,
                    sim.carbon.adder.fed=sim.carbon.adder.fed.bs$mr,
                    sim.ban.lag.0=sim.ban.lag.0.bs$mr)
  deltas.bs = mapply(FUN=function(s, m) list(calc.deltas(sim.base.bs, list(sim=s, mr=m))),
                     sim.list.bs, mr.list.bs)
  names(deltas.bs) = sim.name.list.bs
  sum.table.bs = mapply(FUN=extract.vars.for.table.for.bootstrap, d=deltas.bs, nm=1:length(sim.name.list.bs))
  sum.table.bs = t(sum.table.bs)
  sum.table.bs
  if (any(abs(na.omit(sum.table.bs[1,]))>1e-6)) stop(paste0('Round ',b, ' Baseline did not match.'))
  sum.table.bs = sum.table.bs[-1,] # drop baseline
  
  # rownames(sum.table.bs) = c('$20 Carbon Adder (2%)', 'Moratorium')
  rownames(sum.table.bs) = c('18.75%RROn','25%RROn','25%RROnOff','$20 Carbon Adder (2%)', '$50 Carbon Adder (2%)','Moratorium')
  bootstrap.results[,,b] = sum.table.bs
  pb$tick()
}
dimnames(bootstrap.results)[1:2] = dimnames(sum.table.bs)
save(list='bootstrap.results', file=working.data.dir%&%'/bootstrap_results_'%&%today()%&%'.RData')
bootstrap.results[,,1:5]
ed = Sys.time()
difftime(ed, st) # For 3 scenarios, it took 8 minutes for 3 iterations. 17 hours for 500. Also 17 hours for 201 with 7 scenarios
# load(working.data.loc)

# sum.table.num = read.csv(output.dir%&%subfolder%&%'Summary_table_2021-11-01.csv')
# colnames(sum.table.num) = gsub('\\.',' ', colnames(sum.table.num))

var.temp = 'Gas price'
# $20/tCO2e
dimnames(bootstrap.results)[[1]]
scen.to.plot = 6

dev.off()
par(mfrow=c(3,3))
par(mar=par('mar')*0.7)
plot.new(); text(x=1/2,y=1/2, labels=dimnames(bootstrap.results)[[1]][scen.to.plot], cex=2)
legend('bottom', legend=c('Point Estimate','Mean of Samples'), lty=1:2, lwd=3, col=3)
for (var.temp in dimnames(bootstrap.results)[[2]]) {
  hist(bootstrap.results[scen.to.plot,var.temp,], main=var.temp, freq=F, breaks=30)
  # hist(bootstrap.results[1,var.temp,1:200], main=var.temp, add=TRUE, col=alpha(5, 0.25), freq=F, breaks=30)
  abline(v=c(mean(bootstrap.results[scen.to.plot,var.temp,]), sum.table.num[scen.to.plot+1,var.temp]), lty=2:1, col=3, lwd=3)
  sum.table.num[scen.to.plot+1, var.temp]  
}
apply(bootstrap.results, 1:2, mean)
apply(bootstrap.results, 1:2, mean)/sum.table.num[-1,-1] - 1
apply(bootstrap.results, 1:2, sd)
(apply(bootstrap.results, 1:2, mean)-1.96*apply(bootstrap.results, 1:2, sd))/
apply(bootstrap.results, 1:2, quantile, probs=c(0.025))
(apply(bootstrap.results, 1:2, mean)+1.96*apply(bootstrap.results, 1:2, sd))/
apply(bootstrap.results, 1:2, quantile, probs=c(0.975))
x = apply(bootstrap.results, 1:2, sd)
x[,1:2] = x[,1:2]*100
round(x,3)
round(sum.table.num[-1,-1], 2)

dimnames(bootstrap.results)
dev.off()
pdf(output.dir%&%'/montecarlo_densities_'%&%Sys.Date()%&%'.pdf', family='serif', width=11*1.4, height=8.5)
par(mfrow=c(2,4))
par(mar=c(6.1, 4.1, 4.1, 2.1))
# par(mgp=c(3.5,1,0))
cex.scl = 1.5
for (scen.to.plot in c(4,6)) {
  # plot.new(); par(xpd=NA); text(x=0.4, y=0.5, labels=ifelse(scen.to.plot==4,'$20/tCO2e Carbon\nAdder','Leasing Ban'), cex=2); par(xpd=FALSE)
  # Plot price impacts
  plot(density(100*bootstrap.results[scen.to.plot,'Gas price',], bw=100*0.0005), col=4, main='Price Impacts', lwd=2, xlim=100*c(0,0.025), xlab='', cex=cex.scl, cex.lab=cex.scl, cex.axis=cex.scl, cex.main=cex.scl*1.2)
  title(xlab='Change in Oil or Gas Price (%)', mgp=c(3,1,0), cex.lab=cex.scl)
  grid(); abline(v=0, col=alpha(1, 0.5))
  lines(density(100*bootstrap.results[scen.to.plot,'Oil price',], bw=100*.0005), col=1, lwd=2)
  if (scen.to.plot==4) legend('topright', lty=1, lwd=2, col=c(1,4), legend=c('Oil Price','Gas Price'), bg='white', cex=cex.scl)
  # Plot Delta Emissions
  plot(density(bootstrap.results[scen.to.plot,'Fed US emis',]), col=my.cols[2], main='Emissions Impacts', lwd=2, xlim=c(-350,250), xlab='', ylim=c(0,0.08), cex=cex.scl,cex.lab=cex.scl, cex.axis=cex.scl, cex.main=cex.scl*1.2)
  grid()
  title(xlab='\nChange in Emissions\n(MMTCO2e/year, 2020-2050)', mgp=c(4,1,0), cex.lab=cex.scl)
  lines(density(bootstrap.results[scen.to.plot,'Nonfed US emis',]+bootstrap.results[scen.to.plot,'ROW emis',]), col=my.cols[1], lwd=2)
  lines(density(bootstrap.results[scen.to.plot,'Global emis',], bw=5), col=3, lwd=2) # use ~same bw as others
  abline(v=0, col=alpha(1, 0.5))
  if (scen.to.plot==4) legend('topleft', legend=c('Federal','Nonfederal & ROW','Net Global'), col=c(my.cols[2:1],3), lwd=2, bg='white', cex=cex.scl*1)
  # Leakage rate
  plot(density(100*bootstrap.results[scen.to.plot,'Leakage rate',]), col=1, main='Leakage Rate', lwd=2, xlab='', xlim=c(65,75), cex=cex.scl, cex.lab=cex.scl, cex.axis=cex.scl, cex.main=cex.scl)
  title(xlab='Leakage Rate (%)', mgp=c(3,1,0), cex.lab=cex.scl)
  grid()
  # Revenues
  plot(density(bootstrap.results[scen.to.plot,'Roy and CA Revenue',], bw=0.5), xlim=c(-10,10), col=3, main='Revenue Impacts', lwd=2, xlab='', cex=cex.scl, cex.lab=cex.scl, cex.axis=cex.scl, cex.main=cex.scl*1.2)
  title(xlab='\nChange in Royalty & Carbon\nRevenue ($b/year, 2020-2050)', mgp=c(4,1,0), cex.lab=cex.scl)
  grid(); abline(v=0, col=alpha(1, 0.5))
}
dev.off()
